'''
k折交叉验证代码，
对整体数据集划分五折交叉验证数据，
本章代码可插入到训练代码中，进行五折交叉验证

'''


from sklearn.model_selection import StratifiedKFold
import pandas as pd
import numpy as np
import csv
import os


# 读取整个数据集
def read_data():
    path = 'D:/lung_cancer/data/all_data6.csv'
    data = []
    f = csv.reader(open(path, 'r'))
    for i in f:
        data.append(i)
    return data


# 在所有数据中分离出标签
def divide_label(all_data):
    features = []
    labels = []
    for i in range(1, len(all_data)):
        features.append(all_data[i])
        labels.append(all_data[i][6])

    features = np.asarray(features)
    labels = np.asarray(labels)
    return features, labels


# 依照标签比例分层采样：生成五折数据集
def kfold():
    all_data = read_data()  # 第一行是列名

    features, labels = divide_label(all_data)

    # print(len(features))
    # print(len(labels))

    # 把所有的测试集保存下来，就是相互独立的五折数据集
    sfolder = StratifiedKFold(n_splits=5, random_state=0)

    # 定义五折数据集顺序和csv文件列名
    index = 0
    row_name = all_data[0]

    # 生成存储文件夹
    save_path = 'D:/lung_cancer/data/kfold_dataset'
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    for train, test in sfolder.split(features, labels):
        print('train: %d | test: %d' %(len(train), len(test)))
        son_features = features[test]
        df = pd.DataFrame(son_features, columns=row_name)
        df.to_csv(save_path+'/dataset'+str(index)+'.csv', index=False)
        index = index+1


if __name__ == '__main__':
    kfold()